import cv2
import numpy as np
import os
from glob import glob
import dicom


class Image(object):
    def __init__(self, filename, bw_threshold=128):
        self.data = cv2.imread(filename)
        self.bw_threshold = bw_threshold

    def blockSetGenerator(self, block_height, block_width, height_step, width_step):
        # 一个generator，返回的格式：start行，start列，行数，列数，block的图像
        set_img = self.getSetImg()
        width = self.getWidth()
        height = self.getHeight()
        for x in range(0, height - block_height + 1, height_step):
            for y in range(0, width - block_width + 1, width_step):
                block = set_img[x:x + block_height, y:y + block_width]
                yield x, y, block_height, block_width, block

    def getColorImg(self):
        return self.data

    def getGrayImg(self):
        gray_img = cv2.cvtColor(self.getColorImg(), cv2.COLOR_RGB2GRAY)
        return gray_img

    def getBWImg(self):
        # 0, 255，0表示黑色，255表示白色
        (_, bw_img) = cv2.threshold(self.getGrayImg(), self.bw_threshold, 255, cv2.THRESH_BINARY)
        return bw_img

    def getSetImg(self):
        set_img = (self.getBWImg() / 255).astype(np.int32)
        return set_img

    def getWidth(self):
        return self.data.shape[1]

    def getHeight(self):
        return self.data.shape[0]


class DicomImage(Image):
    def __init__(self, filename, bw_hu_threshold=0):
        self.data = dicom.read_file(filename)
        self.bw_hu_threshold = bw_hu_threshold
        self.hu_image = self.get_pixels_hu()

    def get_pixels_hu(self):
        image = self.data.pixel_array
        image = image.astype(np.int16)
        image[image == -2000] = 0
        intercept = self.data.RescaleIntercept
        slope = self.data.RescaleSlope

        if slope != 1:
            image = slope * image.astype(np.float64)
            image = image.astype(np.int16)
        image += np.int16(intercept)

        return np.array(image, dtype=np.int16)

    def getGrayImg(self):
        raise Exception("not support!")

    def getColorImg(self):
        raise Exception("not support!")

    def getBWImg(self):
        # 0, 255，0表示黑色，255表示白色
        bw_img = self.hu_image.copy()
        bw_img[bw_img <= 0] = 0
        bw_img[bw_img > 0] = 255
        return bw_img

    def getWidth(self):
        return self.data.Columns

    def getHeight(self):
        return self.data.Rows


def glob_all(dir_path, ext='*.jpg'):
    pic_list = glob(os.path.join(dir_path, ext))
    inside = os.listdir(dir_path)
    for dir_name in inside:
        if os.path.isdir(os.path.join(dir_path, dir_name)):
            pic_list.extend(glob_all(os.path.join(dir_path, dir_name), ext))
    return pic_list
